
from vtkmodules.vtkRenderingOpenGL2 import vtkOpenGLActor, vtkOpenGLGlyph3DMapper

from vtkmodules.vtkCommonDataModel import (
    vtkPolyData,
    vtkCellArray,
)

from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkGlyph3DMapper,
    vtkPolyDataMapper,
    vtkShaderProperty,
    vtkActorCollection
)

from vtkmodules.vtkCommonCore import (
    vtkPoints,
    vtkFloatArray,
    vtkUnsignedIntArray,
    vtkUnsignedCharArray,
    vtkBitArray   
)


# instanced rendering 
class InstanceActor(vtkOpenGLActor):
    def __init__(self) -> None:
        super().__init__()
        
        # glyph 3d mapper for instance
        self.mapper = vtkOpenGLGlyph3DMapper()
        self.SetMapper(self.mapper)
        
        # instance data...
        instanceData = vtkPolyData()
        self.mapper.SetInputData(instanceData)  

        self.mapper.SetScalarModeToUsePointFieldData()
        self.mapper.SourceIndexingOn()
        self.mapper.SetSourceIndexArray("SourceIndex")
        self.mapper.SetScaleArray("Scale")
        self.mapper.SetOrientationArray("Orientation")
        self.mapper.SetOrientationModeToRotation()  
        self.mapper.UseSelectionIdsOn()
        self.mapper.SetSelectionIdArray("SelectionId")
        # self.mapper.SelectColorArray("Color")
        # self.mapper.SetColorModeToDirectScalars()        
        self.mapper.MaskingOn()
        self.mapper.SetMaskArray("Masks")

        # instance data
        points = vtkPoints()
        scaleFactors = vtkFloatArray()
        scaleFactors.SetNumberOfComponents(3)
        scaleFactors.SetName("Scale")
        orientations = vtkFloatArray()
        orientations.SetNumberOfComponents(3)
        orientations.SetName("Orientation")
        sourceIndices = vtkUnsignedIntArray()
        sourceIndices.SetName("SourceIndex")
        sourceIndices.SetNumberOfComponents(1)        
        selectionIds = vtkUnsignedIntArray()
        selectionIds.SetName("SelectionId")
        selectionIds.SetNumberOfComponents(1)
        # colors = vtkUnsignedCharArray()
        # #colors.SetName("Color")
        # colors.SetNumberOfComponents(4)
        masks = vtkBitArray()
        masks.SetName("Masks")
        masks.SetNumberOfComponents(1)

        instanceData.SetPoints(points)
        instanceData.GetPointData().AddArray(scaleFactors)
        instanceData.GetPointData().AddArray(orientations)  
        instanceData.GetPointData().AddArray(sourceIndices) 
        instanceData.GetPointData().AddArray(selectionIds) 
        instanceData.GetPointData().AddArray(masks) 
        #instanceData.GetPointData().AddArray(colors)
                        
        self.points = points
        self.scaleFactors = scaleFactors
        self.orientations = orientations  
        self.sourceIndices = sourceIndices
        self.selectionIds = selectionIds
        self.masks = masks
        #self.instanceColors = colors

        # store primitives
        self.primitives = []
        self.instances  = []

    def AddPrimitive(self, prim, node):

        if prim not in self.primitives:
            self.primitives.append(prim)
            self.mapper.SetSourceData(len(self.primitives)-1, prim.GetOutput())             
        # polyData index
        primIdx = self.primitives.index(prim)

        # transform
        #tm = node.worldTransform
        # initial values...
        scale = [1, 1, 1]
        position = [0, 0, 0]
        orientation = [0, 0, 0]
        node.GetPosition(position)
        node.GetScale(scale)
        node.GetOrientation(orientation)
        self.points.InsertNextPoint(position[0], position[1], position[2])
        self.scaleFactors.InsertNextTuple3(scale[0], scale[1], scale[2])
        self.orientations.InsertNextTuple3(orientation[0], orientation[1], orientation[2])
        self.sourceIndices.InsertNextValue(primIdx)        
        self.selectionIds.InsertNextValue(len(self.instances))          
        self.masks.InsertNextValue(node.visible)
        #self.instanceColors.InsertNextTuple4(255, 255, 255, 255)  
        
        self.instances.append((primIdx, node))

    def Update(self):
        
        changed = False
        # update instance data
        for i in range(len(self.instances)):

            idx, node = self.instances[i]
            #tm = node.worldTransform
            #if tm.GetMTime() > self.GetMTime() or node.GetMTime() > self.GetMTime():
            if node.GetMTime() > self.GetMTime():
                scale = [1, 1, 1]
                position = [0, 0, 0]
                orientation = [0, 0, 0]
                node.GetPosition(position)
                node.GetScale(scale)
                node.GetOrientation(orientation)
                self.points.SetPoint(i, position[0], position[1], position[2])
                self.scaleFactors.SetTuple3(i, scale[0], scale[1], scale[2])
                self.orientations.SetTuple3(i, orientation[0], orientation[1], orientation[2])
                    
                # visibility
                self.masks.SetValue(i, node.visible)
                #self.sourceIndices.SetValue(i, idx)
                changed = True        
        if changed:
            self.mapper.Modified()                


    def GetNode(self, idx):
        return self.instances[idx][1]
